"""
激活值钩子模块，用于监控模型层的激活值。
"""
from typing import Callable, Optional
import torch

class ActivationHook:
    """激活值钩子"""
    
    def __init__(self,
                 name: str,
                 module: torch.nn.Module,
                 callback: Callable[[str, torch.Tensor, bool], None]):
        """
        初始化激活值钩子
        
        Args:
            name: 层名称
            module: 层模块
            callback: 回调函数
        """
        self.name = name
        self.module = module
        self.callback = callback
        self.hook = None
        self.register()
    
    def register(self) -> None:
        """注册钩子"""
        if self.hook is not None:
            return
        
        def hook(module: torch.nn.Module, inputs: tuple, output: torch.Tensor) -> None:
            """激活值钩子函数"""
            # 检查是否在训练模式
            is_training = module.training
            
            # 调用回调函数
            self.callback(self.name, output, is_training)
        
        # 注册钩子
        self.hook = self.module.register_forward_hook(hook)
    
    def remove(self) -> None:
        """移除钩子"""
        if self.hook is not None:
            self.hook.remove()
            self.hook = None
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.remove() 